# src/evaluate.py

"""
模型评估脚本 - 支持单任务和多任务评估

用法：
    python src/evaluate.py --checkpoint models/best_model_age.pth --split test --task age
    python src/evaluate.py --checkpoint models/best_model_gender.pth --split test --task gender
    python src/evaluate.py --checkpoint models/best_model_both.pth --split test --task both
"""

import sys
import os
# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import torch.nn as nn
import numpy as np
from tqdm import tqdm
from sklearn.metrics import (
    accuracy_score, f1_score, roc_auc_score, 
    confusion_matrix, classification_report,
    mean_absolute_error, mean_squared_error, r2_score, precision_score, recall_score
)
import matplotlib.pyplot as plt
import seaborn as sns
import argparse

from src.config_loader import load_config
from src.dataset import HRVDataset
from src.model import ChannelIndependentTransformer, HRVHybridTransformer


def evaluate_model(model, data_loader, device, task_type='both', loss_alpha=0.5, verbose=True, use_multimodal=False):
    """
    评估模型性能（支持单任务和多任务）
    
    Args:
        model: 要评估的模型
        data_loader: 数据加载器
        device: 运行设备
        task_type: 任务类型 ('age', 'gender', 或 'both')
        loss_alpha: 损失函数权重（仅在task_type='both'时使用）
        verbose: 是否打印详细信息
        use_multimodal: 是否使用多模态输入
        
    Returns:
        dict: 包含所有评估指标的字典
    """
    model.eval()
    
    all_age_preds, all_ages = [], []
    all_gender_probs, all_genders = [], []
    
    # 用于计算loss
    if task_type in ['age', 'both']:
        loss_fn_age = nn.MSELoss()
    if task_type in ['gender', 'both']:
        loss_fn_gender = nn.BCELoss()
    
    total_loss = 0.0
    
    with torch.no_grad():
        iterator = tqdm(data_loader, desc="Evaluating") if verbose else data_loader
        
        for batch in iterator:
            # 根据是否使用多模态解包数据
            if use_multimodal:
                windows, tabular_features, ages, genders = batch
                tabular_features = tabular_features.to(device)
            else:
                windows, ages, genders = batch
                tabular_features = None
            
            windows = windows.to(device)
            ages = ages.to(device).squeeze(-1)
            genders = genders.to(device).squeeze(-1).float()
            
            # 前向传播（根据任务类型）
            if task_type == 'age':
                if use_multimodal:
                    age_preds = model(windows, tabular_features)
                else:
                    age_preds = model(windows)
                loss = loss_fn_age(age_preds, ages)
                total_loss += loss.item()
                all_age_preds.extend(age_preds.cpu().numpy())
                all_ages.extend(ages.cpu().numpy())
                
            elif task_type == 'gender':
                if use_multimodal:
                    gender_logits = model(windows, tabular_features)
                else:
                    gender_logits = model(windows)
                gender_probs = torch.sigmoid(gender_logits)
                loss = loss_fn_gender(gender_probs, genders)
                total_loss += loss.item()
                all_gender_probs.extend(gender_probs.cpu().numpy())
                all_genders.extend(genders.cpu().numpy())
                
            else:  # 'both'
                if use_multimodal:
                    age_preds, gender_logits = model(windows, tabular_features)
                else:
                    age_preds, gender_logits = model(windows)
                gender_probs = torch.sigmoid(gender_logits)
                
                loss_age = loss_fn_age(age_preds, ages)
                loss_gender = loss_fn_gender(gender_probs, genders)
                combined_loss = loss_alpha * loss_age + (1 - loss_alpha) * loss_gender
                total_loss += combined_loss.item()
                
                all_age_preds.extend(age_preds.cpu().numpy())
                all_ages.extend(ages.cpu().numpy())
                all_gender_probs.extend(gender_probs.cpu().numpy())
                all_genders.extend(genders.cpu().numpy())
    
    # 计算指标
    metrics = {}
    metrics['loss'] = total_loss / len(data_loader)
    
    # 年龄回归指标
    if task_type in ['age', 'both']:
        age_preds_arr = np.array(all_age_preds)
        ages_arr = np.array(all_ages)
        
        metrics['age_mae'] = mean_absolute_error(ages_arr, age_preds_arr)
        metrics['age_rmse'] = np.sqrt(mean_squared_error(ages_arr, age_preds_arr))
        metrics['age_r2'] = r2_score(ages_arr, age_preds_arr)
        
        metrics['age_predictions'] = {
            'true': ages_arr,
            'pred': age_preds_arr
        }
    
    # 性别分类指标
    if task_type in ['gender', 'both']:
        gender_probs_arr = np.array(all_gender_probs)
        genders_arr = np.array(all_genders)
        gender_preds_arr = (gender_probs_arr > 0.5).astype(int)
        
        metrics['gender_accuracy'] = accuracy_score(genders_arr, gender_preds_arr)
        metrics['gender_precision'] = precision_score(genders_arr, gender_preds_arr, zero_division=0)
        metrics['gender_recall'] = recall_score(genders_arr, gender_preds_arr, zero_division=0)
        metrics['gender_f1'] = f1_score(genders_arr, gender_preds_arr, zero_division=0)
        
        try:
            metrics['gender_auc'] = roc_auc_score(genders_arr, gender_probs_arr)
        except ValueError:
            metrics['gender_auc'] = 0.5  # 只有一个类别时无法计算AUC
        
        metrics['gender_confusion_matrix'] = confusion_matrix(genders_arr, gender_preds_arr)
        metrics['gender_classification_report'] = classification_report(
            genders_arr, gender_preds_arr, 
            target_names=['女', '男'],
            zero_division=0
        )
        
        metrics['gender_predictions'] = {
            'true': genders_arr,
            'prob': gender_probs_arr,
            'pred': gender_preds_arr
        }
    
    return metrics


def print_metrics(metrics, task_type='both', split_name="Test"):
    """
    打印评估指标
    
    Args:
        metrics: 评估指标字典
        task_type: 任务类型
        split_name: 数据集名称
    """
    print("\n" + "="*70)
    print(f"{split_name} 集评估结果 - 任务: {task_type.upper()}")
    print("="*70)
    
    print(f"\n📊 总体Loss: {metrics['loss']:.4f}")
    
    # 年龄预测指标
    if task_type in ['age', 'both']:
        print(f"\n👤 年龄预测指标:")
        print(f"  - MAE (平均绝对误差):  {metrics['age_mae']:.2f} 岁")
        print(f"  - RMSE (均方根误差):    {metrics['age_rmse']:.2f} 岁")
        print(f"  - R² Score:            {metrics['age_r2']:.4f}")
    
    # 性别分类指标
    if task_type in ['gender', 'both']:
        print(f"\n🚻 性别分类指标:")
        print(f"  - 准确率 (Accuracy):    {metrics['gender_accuracy']:.4f} ({metrics['gender_accuracy']*100:.2f}%)")
        print(f"  - 精确率 (Precision):   {metrics['gender_precision']:.4f}")
        print(f"  - 召回率 (Recall):      {metrics['gender_recall']:.4f}")
        print(f"  - F1-Score:            {metrics['gender_f1']:.4f}")
        print(f"  - AUC-ROC:             {metrics['gender_auc']:.4f}")
        
        print(f"\n混淆矩阵:")
        cm = metrics['gender_confusion_matrix']
        print(f"         女    男")
        print(f"  女   [{cm[0,0]:4d} {cm[0,1]:4d}]")
        print(f"  男   [{cm[1,0]:4d} {cm[1,1]:4d}]")
        
        print(f"\n分类报告:")
        print(metrics['gender_classification_report'])


def plot_age_results(metrics, save_path=None):
    """
    绘制年龄预测结果
    """
    age_true = metrics['age_predictions']['true']
    age_pred = metrics['age_predictions']['pred']
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # 散点图
    ax1 = axes[0]
    ax1.scatter(age_true, age_pred, alpha=0.5, s=30)
    ax1.plot([age_true.min(), age_true.max()], 
             [age_true.min(), age_true.max()], 
             'r--', lw=2, label='Perfect Prediction')
    ax1.set_xlabel('True Age (years)', fontsize=12)
    ax1.set_ylabel('Predicted Age (years)', fontsize=12)
    ax1.set_title(f'Age Prediction\nMAE={metrics["age_mae"]:.2f}, R²={metrics["age_r2"]:.3f}', 
                  fontsize=14, fontweight='bold')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # 误差分布
    ax2 = axes[1]
    errors = age_pred - age_true
    ax2.hist(errors, bins=30, edgecolor='black', alpha=0.7)
    ax2.axvline(x=0, color='r', linestyle='--', linewidth=2, label='Zero Error')
    ax2.set_xlabel('Prediction Error (years)', fontsize=12)
    ax2.set_ylabel('Frequency', fontsize=12)
    ax2.set_title(f'Age Prediction Error Distribution\nRMSE={metrics["age_rmse"]:.2f}', 
                  fontsize=14, fontweight='bold')
    ax2.legend()
    ax2.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"✅ 年龄预测图表已保存: {save_path}")
    else:
        plt.show()
    
    plt.close()


def plot_gender_results(metrics, save_path=None):
    """
    绘制性别分类结果
    """
    from sklearn.metrics import roc_curve
    
    gender_true = metrics['gender_predictions']['true']
    gender_prob = metrics['gender_predictions']['prob']
    cm = metrics['gender_confusion_matrix']
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # 混淆矩阵
    ax1 = axes[0]
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=['女', '男'], yticklabels=['女', '男'],
                ax=ax1, cbar_kws={'label': 'Count'})
    ax1.set_xlabel('Predicted', fontsize=12)
    ax1.set_ylabel('True', fontsize=12)
    ax1.set_title(f'Gender Classification Confusion Matrix\nAccuracy={metrics["gender_accuracy"]:.3f}, F1={metrics["gender_f1"]:.3f}', 
                  fontsize=14, fontweight='bold')
    
    # ROC曲线
    ax2 = axes[1]
    if metrics['gender_auc'] > 0.5:  # 确保可以计算ROC
        fpr, tpr, _ = roc_curve(gender_true, gender_prob)
        ax2.plot(fpr, tpr, linewidth=2, label=f'ROC Curve (AUC = {metrics["gender_auc"]:.3f})')
        ax2.plot([0, 1], [0, 1], 'k--', linewidth=2, label='Random Guess')
        ax2.set_xlabel('False Positive Rate', fontsize=12)
        ax2.set_ylabel('True Positive Rate', fontsize=12)
        ax2.set_title('ROC Curve', fontsize=14, fontweight='bold')
        ax2.legend(loc='lower right')
        ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"✅ 性别分类图表已保存: {save_path}")
    else:
        plt.show()
    
    plt.close()


def main():
    """
    主评估函数
    """
    parser = argparse.ArgumentParser(description='评估HRV Transformer模型')
    parser.add_argument('--checkpoint', type=str, default=None, 
                       help='模型checkpoint路径')
    parser.add_argument('--split', type=str, default='test', choices=['train', 'val', 'test'],
                       help='评估哪个数据集')
    parser.add_argument('--task', type=str, default=None, choices=['age', 'gender', 'both'],
                       help='任务类型（如果不指定，从config.yaml读取）')
    parser.add_argument('--no_plot', action='store_true',
                       help='不生成图表')
    
    args = parser.parse_args()
    
    # 加载配置
    cfg = load_config()
    
    # 确定任务类型
    if args.task is not None:
        task_type = args.task
    else:
        task_type = cfg.get('task', {}).get('type', 'both')
    
    print("\n" + "="*70)
    print(f"模型评估 - 任务: {task_type.upper()}")
    print("="*70)
    
    # 确定checkpoint路径
    if args.checkpoint is None:
        model_save_dir = cfg['output']['model_save_dir']
        checkpoint_path = os.path.join(model_save_dir, f"best_model_{task_type}.pth")
    else:
        checkpoint_path = args.checkpoint
    
    if not os.path.exists(checkpoint_path):
        print(f"❌ 错误: checkpoint文件不存在: {checkpoint_path}")
        return
    
    print(f"\n📂 加载checkpoint: {checkpoint_path}")
    
    # 设置设备
    device = torch.device(cfg['device'] if torch.cuda.is_available() else 'cpu')
    print(f"🖥️  使用设备: {device}")
    
    # 确定模型类型和多模态设置
    model_type = cfg.get('model', {}).get('type', 'channel_independent')
    use_multimodal = cfg.get('multimodal', {}).get('enabled', False)
    
    # 加载数据
    print(f"\n加载 {args.split.upper()} 数据集...")
    dataset = HRVDataset(split=args.split, cfg=cfg)
    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=cfg['training']['batch_size'],
        shuffle=False,
        num_workers=cfg['training']['num_workers'],
        pin_memory=True if device.type == 'cuda' else False
    )
    print(f"✅ 数据加载完成: {len(dataset)} 个样本")
    
    # 加载模型
    print(f"\n初始化模型...")
    if model_type == 'hybrid':
        model = HRVHybridTransformer(cfg=cfg, task_type=task_type, use_multimodal=use_multimodal)
        print(f"  模型类型: HRVHybridTransformer (CNN + Transformer混合架构)")
    else:
        model = ChannelIndependentTransformer(cfg=cfg, task_type=task_type)
        print(f"  模型类型: ChannelIndependentTransformer (纯Transformer)")
    
    print(f"  多模态: {'启用' if use_multimodal else '禁用'}")
    
    model.load_state_dict(torch.load(checkpoint_path, map_location=device))
    model = model.to(device)
    print(f"✅ 模型加载完成")
    
    # 评估
    print(f"\n开始评估...")
    loss_alpha = cfg['training']['loss_alpha'] if task_type == 'both' else None
    metrics = evaluate_model(
        model=model,
        data_loader=data_loader,
        device=device,
        task_type=task_type,
        loss_alpha=loss_alpha,
        verbose=True,
        use_multimodal=use_multimodal
    )
    
    # 打印指标
    print_metrics(metrics, task_type=task_type, split_name=args.split.upper())
    
    # 生成图表
    if not args.no_plot:
        results_dir = cfg['output']['results_dir']
        os.makedirs(results_dir, exist_ok=True)
        
        print(f"\n生成评估图表...")
        
        if task_type in ['age', 'both']:
            age_plot_path = os.path.join(results_dir, f'evaluation_{args.split}_age.png')
            plot_age_results(metrics, save_path=age_plot_path)
        
        if task_type in ['gender', 'both']:
            gender_plot_path = os.path.join(results_dir, f'evaluation_{args.split}_gender.png')
            plot_gender_results(metrics, save_path=gender_plot_path)
        
        print(f"\n✅ 所有图表已保存到: {results_dir}/")
    
    print("\n" + "="*70)
    print("评估完成！")
    print("="*70)


if __name__ == "__main__":
    main()
